import os
from typing import List, Dict, Any
from pypdf import PdfReader
import pytesseract
from pdf2image import convert_from_path
from PIL import Image
import cv2
import numpy as np
from langchain.text_splitter import RecursiveCharacterTextSplitter

class PDFProcessor:
    def __init__(self, pdf_path: str):
        self.pdf_path = pdf_path
        self.reader = PdfReader(pdf_path)
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len
        )
        # 配置Tesseract
        self.tesseract_config = '--oem 3 --psm 6 -l chi_sim+eng'  # 使用中文和英文语言包

    def preprocess_image(self, image: np.ndarray) -> np.ndarray:
        """图像预处理以提高OCR准确率"""
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 自适应阈值处理
        binary = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY, 11, 2
        )
        
        # 降噪
        denoised = cv2.fastNlMeansDenoising(binary)
        
        # 锐化
        kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]])
        sharpened = cv2.filter2D(denoised, -1, kernel)
        
        return sharpened

    def extract_text(self) -> List[str]:
        """提取PDF中的文本内容"""
        text_chunks = []
        for page in self.reader.pages:
            text = page.extract_text()
            if text:
                chunks = self.text_splitter.split_text(text)
                text_chunks.extend(chunks)
        return text_chunks

    def extract_images(self) -> List[Dict[str, Any]]:
        """提取PDF中的图片并进行OCR处理"""
        images = []
        pdf_images = convert_from_path(
            self.pdf_path,
            dpi=300,  # 提高DPI以获得更好的图像质量
            thread_count=4  # 多线程处理
        )
        
        for page_num, image in enumerate(pdf_images):
            # 转换为OpenCV格式
            opencv_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
            
            # 图像预处理
            processed_image = self.preprocess_image(opencv_image)
            
            # 使用OCR提取图片中的文字
            text = pytesseract.image_to_string(
                processed_image,
                config=self.tesseract_config
            )
            
            # 如果识别结果为空，尝试使用原始图像
            if not text.strip():
                text = pytesseract.image_to_string(
                    opencv_image,
                    config=self.tesseract_config
                )
            
            # 将图片转换为base64格式
            _, buffer = cv2.imencode('.jpg', opencv_image)
            image_base64 = buffer.tobytes()
            
            # 添加置信度信息
            confidence = pytesseract.image_to_data(
                processed_image,
                config=self.tesseract_config,
                output_type=pytesseract.Output.DICT
            )
            
            # 计算平均置信度
            confidences = [int(conf) for conf in confidence['conf'] if conf != '-1']
            avg_confidence = sum(confidences) / len(confidences) if confidences else 0
            
            images.append({
                'page': page_num + 1,
                'image': image_base64,
                'text': text,
                'confidence': avg_confidence
            })
        
        return images

    def process_document(self) -> Dict[str, Any]:
        """处理整个PDF文档，返回文本和图片信息"""
        return {
            'text_chunks': self.extract_text(),
            'images': self.extract_images()
        } 